package app.jppfnetmorph;

import static optimization.LinearAlgebra.vAdd;
import static optimization.LinearAlgebra.vMul;
import static optimization.LinearAlgebra.vNorm2;
import static optimization.LinearAlgebra.vPrint;
import static optimization.LinearAlgebra.vSelect;
import static optimization.LinearAlgebra.vSub;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Random;

import neuron.NetworkStatisticsProxy.Selector;
import optimization.LinearAlgebra.VectorFunction;
import netmorph.model.parts.NetmorphModelBESPart;
import netmorph.model.parts.NetmorphModelElongPart;

import org.jppf.client.JPPFJob;

import common.statistics.Series;


public class NelderMeadParallel {
	
	class PointValue implements Comparable<PointValue> {
		double[] x;
		double fx;
		public PointValue(double[] p) {
			this.x = p;
		}
		@Override
		public int compareTo(PointValue o) {
			return Double.compare(fx, o.fx);
		}
	}

	/**
	 * Simplex points
	 */
	List<PointValue> points;
	

	/**
	 * Constants
	 */
	protected double alpha = 1, beta = 0.5, gamma = 2, delta = 0.9;
	

	/**
	 * Constraints mapping
	 */
	VectorFunction cons;
	
	/**
	 * Dennis-Woods termination threshold
	 */
	double nu = 0.0001;
	
	Selector sel;
	Class<? extends NetmorphModelBESPart> partClass;
	NetmorphConfig.Initializer init;
	
	double[] yt;
	int samples;
	
	public NelderMeadParallel(Selector sel, Class<? extends NetmorphModelBESPart> partClass, NetmorphConfig.Initializer init, int samples, double[] yt, double nu, VectorFunction cons) {
		super();
		this.sel = sel;
		this.partClass = partClass;
		this.init = init;
		this.samples = samples;
		this.yt = yt.clone();
		this.cons = cons;
		this.nu = nu;
	}

	private void setP0(double[] nx, double nfx)
	{
		points.get(0).x = nx;
		points.get(0).fx = nfx;
	}

	private void setP1(double[] nx, double nfx)
	{
		points.get(1).x = nx;
		points.get(1).fx = nfx;
	}

	private void setP2(double[] nx, double nfx)
	{
		points.get(2).x = nx;
		points.get(2).fx = nfx;
	}



	
	protected boolean dennisWoodsTerminationTest()
	{
		double normX0 = Math.max(1,vNorm2(points.get(0).x));
		double maxNorm = Double.MIN_VALUE;
		for (int i = 1; i < 3; i++) {
			double norm = vNorm2(vSub(points.get(0).x, points.get(i).x));
			if (norm > maxNorm) maxNorm = norm;
		}
		if (maxNorm / normX0 < nu) {
			System.out.println("NM-TERM-DENNISWOODS " + (maxNorm / normX0) + " < " + nu);
			return true;
		}
		return false;
	}
	
	
	public static JPPFJob compJob(List<double[]> points, int N, String jobID, Class<? extends NetmorphModelBESPart> partClass, NetmorphConfig.Initializer init) throws Exception
	{
		return compJob(points, N, jobID, partClass, null, init);
	}
	
	public static JPPFJob compJob(List<double[]> points, int N, String jobID, Class<? extends NetmorphModelBESPart> partClass, Class<? extends NetmorphModelElongPart> elongClass, NetmorphConfig.Initializer init) throws Exception
	{
		JPPFJob job = new JPPFJob();
		job.setId(jobID);
		
		Random rnd = new Random();
		
		int pointIter = 0;
		
		for (double[] x : points) {
			int pop = 10;
			int times = (int) Math.ceil(1.0*N / (pop*7));
			pointIter++;
			
			NetmorphConfig cfg = new NetmorphConfig(init);
			cfg.setNeuronPopulationSize(pop);
			
			for (int i = 0; i < times; i++) {
				long seed = rnd.nextLong();
				String config = cfg.createConfig(x, seed, partClass, elongClass);
				String id = "ParNM-" + jobID + "-" + pointIter + "-" + i;
				JPPFNetmorphTask task = new JPPFNetmorphTask(config, id, x, i);
				job.addTask(task);
//				System.out.println("Added task " + id);
			}
		}

		return job;
	}
	

	public JPPFJob initJob(int jobN) throws Exception
	{
		List<double[]> lst = new ArrayList<double[]>();
		for (PointValue pv : points) {
			lst.add(pv.x);
		}
		return compJob(lst, jobN, "Init", partClass, init);
	}

	
	public JPPFJob stepJob(int jobN, double[] xr, double[] xe, double[] xcout, double[] xcin) throws Exception
	{
		List<double[]> lst = new ArrayList<double[]>();
		lst.add(xr);
		lst.add(xe);
		lst.add(xcout);
		lst.add(xcin);
		return compJob(lst, jobN, "Step", partClass, init);
	}

	
	private void shrink() throws Exception
	{
		double[] xs0 = points.get(0).x;
		double[] xs1 = cons.evaluate(vAdd(points.get(0).x, vMul(getDelta(), vSub(points.get(1).x, points.get(0).x))));
		double[] xs2 = cons.evaluate(vAdd(points.get(0).x, vMul(getDelta(), vSub(points.get(2).x, points.get(0).x))));

//		System.out.println("NMP shrink run");		
		List<JobResult> jrl = JPPFRunner.runJob(shrinkJob(samples, xs0, xs1, xs2));
//		System.out.println("NMP shrink parse");		

		double fs0 = eval(jrl, xs0, yt, sel);
		double fs1 = eval(jrl, xs1, yt, sel);
		double fs2 = eval(jrl, xs2, yt, sel);

		setP0(xs0, fs0);
		setP1(xs1, fs1);
		setP2(xs2, fs2);
	}
		

	public JPPFJob shrinkJob(int jobN, double[] xs0, double[] xs1, double[] xs2) throws Exception
	{
		List<double[]> lst = new ArrayList<double[]>();
		lst.add(xs0);
		lst.add(xs1);
		lst.add(xs2);
		return compJob(lst, jobN, "Shrink", partClass, init);
	}

	
	public static Series collect(List<JobResult> res, double[] x, Selector sel)
	{
		Series s = new Series();
		for (JobResult jr : res) {
			// if result is for this x
			if (jr.isX(x)) {
				// merge requested series 
				s.merge(sel.get(jr.stats));
			}
		}
		return s;
	}
	
	
	public static double eval(List<JobResult> res, double[] x, double[] yt, Selector sel)
	{
		Series fx = collect(res, x, sel);
		return vNorm2(vSub(fx.meanAndDev(), yt));
	}

	
	private void init(double[] x0i, double[] h0) throws Exception
	{
		points = new ArrayList<NelderMeadParallel.PointValue>();
		
		points.add(new PointValue(x0i));
		points.add(new PointValue(vAdd(x0i, vSelect(h0, 0))));
		points.add(new PointValue(vAdd(x0i, vSelect(h0, 1))));
		
//		System.out.println("NMP init run");
		List<JobResult> jrl = JPPFRunner.runJob(initJob(samples));
//		System.out.println("NMP init parse");
		points.get(0).fx = eval(jrl, points.get(0).x, yt, sel);
		points.get(1).fx = eval(jrl, points.get(1).x, yt, sel);
		points.get(2).fx = eval(jrl, points.get(2).x, yt, sel);
	}

	
	private void step() throws Exception
	{
		double fl, fs, fh;
		fl = points.get(0).fx; // best        0
		fs = points.get(1).fx; // 2nd worst   n-2
		fh = points.get(2).fx; // worst       n-1

		// compute centroid
		double[] c = vAdd(points.get(0).x, points.get(1).x);
		c = vMul(1.0/2.0, c);
		
		System.out.format("f(x):   %.5f, %.5f, %.5f \n", fl, fs, fh);
		
		
		// compute reflection, expansion and contraction points
		double[] xr = cons.evaluate(vAdd(c, vMul(getAlpha(), vSub(c, points.get(2).x))));
		double[] xe = cons.evaluate(vAdd(c, vMul(getGamma(), vSub(xr, c))));
		double[] xcout = cons.evaluate(vAdd(c, vMul(getBeta(), vSub(xr, c))));
		double[] xh = points.get(2).x;
		double[] xcin = cons.evaluate(vAdd(c, vMul(getBeta(), vSub(xh, c))));

//		System.out.println("NMP step run");		
		List<JobResult> jrl = JPPFRunner.runJob(stepJob(samples, xr, xe, xcout, xcin));
//		System.out.println("NMP step parse");		

		double fr = eval(jrl, xr, yt, sel);
		double fe = eval(jrl, xe, yt, sel);
		double fcout = eval(jrl, xcout, yt, sel);
		double fcin = eval(jrl, xcin, yt, sel);
		

		// STEP ?
		if (fl <= fr && fr < fs) {
			System.out.println("R");
			setP2(xr, fr);
		}
		// EXPAND ?
		else if (fr < fl && fe < fr) {
			System.out.println("E1");
			setP2(xe, fe);
		}
		else if (fr < fl) {
			System.out.println("ER");
			setP2(xr, fr);
		}
		// CONTRACT ?
		else if (fr >= fs && fr < fh && fcout <= fr) {
			System.out.println("CO");
			setP2(xcout, fcout);
		}
		else if (fr >= fs && fr >= fh && fcin < fh) {
			System.out.println("CI");
			setP2(xcin, fcin);
		}
		// SHRINK
		else {
			System.out.println("SS");
			shrink();
		}
	}
	
	public double[] search(double[] x0i, double[] h0, int maxIter) throws Exception 
	{
		// compute initial simplex points
		init(x0i, h0);

		
		for (int iter = 0; iter < maxIter; iter++) {
			
			// sort by value
			Collections.sort(points);

			// termination tests
			if (iter > 3 && dennisWoodsTerminationTest()) {
				System.out.println("NM-" + "   " + vPrint(points.get(0).x)+ ", " + vPrint(points.get(1).x) + ", " + vPrint(points.get(2).x));
				return points.get(0).x;
			}
						
			System.out.print("NM"+iter + "   " + vPrint(points.get(0).x)+ ", " + vPrint(points.get(1).x) + ", " + vPrint(points.get(2).x));
	
			// try reflection, expansion or contraction
			step();			

		} // ITER
		
		System.out.println("NM-TERM-MAXITER");
		return points.get(0).x;
	}


	public double getEstimateError()
	{
		return points.get(0).fx;
	}

	protected double getAlpha() {
		return alpha;
	}

	protected double getBeta() {
		return beta;
	}

	protected double getGamma() {
		return gamma;
	}

	protected double getDelta() {
		return delta;
	}
}
